import os
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
import time

def process_csv_files(base_path, output_path, feature_choice='Voltage'):
    """
    处理CSV文件并创建多分类数据集

    参数:
    base_path: 包含类别文件夹的根路径
    output_path: 输出文件的路径
    feature_choice: 选择的特征('Current'或'Voltage')
    """
    train_path = os.path.join(output_path, 'train')
    test_path = os.path.join(output_path, 'test')
    os.makedirs(train_path, exist_ok=True)
    os.makedirs(test_path, exist_ok=True)

    class_folders = sorted([f for f in os.listdir(base_path)
                            if os.path.isdir(os.path.join(base_path, f))])
    class_mapping = {folder: idx for idx, folder in enumerate(class_folders)}

    # 写入类别标签映射文件
    label_map_path = os.path.join(output_path, 'class_labels.txt')
    with open(label_map_path, 'w', encoding='utf-8') as f:
        for class_name, label in class_mapping.items():
            f.write(f"class{label}: {class_name}\n")
    print(f"\n已生成类别标签映射文件: {label_map_path}")

    file_paths = []
    labels = []

    print("正在扫描文件夹和文件...")

    for class_name in class_folders:
        class_idx = class_mapping[class_name]
        class_dir = os.path.join(base_path, class_name)

        csv_files = [f for f in os.listdir(class_dir) if f.endswith('.csv')]
        print(f"类别 {class_name} (标签 {class_idx}): 找到 {len(csv_files)} 个文件")

        for file_name in csv_files:
            file_path = os.path.join(class_dir, file_name)
            file_paths.append(file_path)
            labels.append(class_idx)

    file_paths = np.array(file_paths)
    labels = np.array(labels)

    print(f"\n总共找到 {len(file_paths)} 个文件")
    for cls, count in zip(*np.unique(labels, return_counts=True)):
        print(f"Class {cls}: {count} 个文件")

    # 划分训练测试集（保持分层抽样）
    train_paths, test_paths, train_labels, test_labels = train_test_split(
        file_paths, labels, test_size=0.2, random_state=42, stratify=labels
    )

    print(f"\n训练集大小: {len(train_paths)}")
    print(f"测试集大小: {len(test_paths)}")
    print("\n处理训练集...")
    process_dataset(train_paths, train_labels, train_path, feature_choice)
    print("\n处理测试集...")
    process_dataset(test_paths, test_labels, test_path, feature_choice)
    print(f"\n处理完成！文件已保存到: {output_path}")

def process_dataset(file_paths, labels, output_base_path, feature_choice):
    """
    处理数据集并保存为NPY格式

    参数:
    file_paths: 文件路径列表
    labels: 对应的标签列表
    output_base_path: 输出根目录（如train/或test/）
    feature_choice: 选择的特征
    """
    class_dirs = {}
    for label in np.unique(labels):
        class_dir = os.path.join(output_base_path, f'class{label}')
        os.makedirs(class_dir, exist_ok=True)
        class_dirs[label] = class_dir

    for idx, (file_path, label) in enumerate(zip(file_paths, labels)):
        try:
            col_index = 0 if feature_choice == 'Current' else 1
            data = pd.read_csv(file_path, usecols=[col_index], skiprows=1, header=None)
            data_array = data.values.flatten()

            file_name = os.path.basename(file_path).replace('.csv', '.npy')
            output_path = os.path.join(class_dirs[label], file_name)

            np.save(output_path, data_array)

        except Exception as e:
            print(f"处理文件 {file_path} 时出错: {str(e)}")

        if (idx + 1) % 100 == 0:
            print(f"已处理 {idx + 1}/{len(file_paths)} 个文件")

def main():
    input_path = r"D:\A_90762\atl\所有NG1"

    output_path = r"D:\A_90762\atl\Current"
    feature = "Current"  # 可选择 "Current" 或 "Voltage"

    # output_path = r"D:\A_90762\atl\Voltage"
    # feature = "Voltage"

    print(f"输入路径: {input_path}")
    print(f"输出路径: {output_path}")
    print(f"选择的特征: {feature}")

    if not os.path.exists(input_path):
        print(f"错误: 输入路径 '{input_path}' 不存在!")
        return

    start_time = time.time()
    process_csv_files(input_path, output_path, feature)
    end_time = time.time()

    print(f"总处理时间: {end_time - start_time:.2f} 秒")

if __name__ == "__main__":
    main()

